#!/bin/bash

# 快速启动优化训练脚本
# 适合熟悉配置的用户快速开始

echo "🚀 快速启动优化训练"
echo ""

# 配置选项
echo "请选择训练配置:"
echo ""
echo "1) 标准优化 (YOLOv8n, 800px, 120epochs) - 推荐"
echo "2) 快速测试 (YOLOv8n, 640px, 50epochs) - 快速验证"
echo "3) 高精度 (YOLOv8n, 800px, 150epochs) - 充分训练"
echo "4) 超高精度 (YOLOv8n, 960px, 200epochs) - 极致训练"
echo "5) 自定义配置"
echo ""

read -p "请输入选项 (1-5): " choice

case $choice in
    1)
        MODEL="yolov8n.pt"
        IMGSZ=800
        EPOCHS=120
        BATCH=16
        NAME="urpc2020_optimized_standard"
        echo "✅ 选择: 标准优化配置"
        ;;
    2)
        MODEL="yolov8n.pt"
        IMGSZ=640
        EPOCHS=50
        BATCH=16
        NAME="urpc2020_optimized_quick"
        echo "✅ 选择: 快速测试配置"
        ;;
    3)
        MODEL="yolov8n.pt"
        IMGSZ=800
        EPOCHS=150
        BATCH=16
        NAME="urpc2020_optimized_high"
        echo "✅ 选择: 高精度配置"
        ;;
    4)
        MODEL="yolov8n.pt"
        IMGSZ=960
        EPOCHS=200
        BATCH=12
        NAME="urpc2020_optimized_ultra"
        echo "✅ 选择: 超高精度配置"
        echo "⚠️  注意: 需要更多显存和训练时间 (~4小时)"
        ;;
    5)
        echo "自定义配置:"
        read -p "  模型 (yolov8n/s/m/l/x): " MODEL
        MODEL="${MODEL}.pt"
        read -p "  图像尺寸 (640/800/960/1024): " IMGSZ
        read -p "  训练轮数: " EPOCHS
        read -p "  批次大小: " BATCH
        read -p "  实验名称: " NAME
        echo "✅ 选择: 自定义配置"
        ;;
    *)
        echo "❌ 无效选项，使用默认配置"
        MODEL="yolov8n.pt"
        IMGSZ=800
        EPOCHS=120
        BATCH=16
        NAME="urpc2020_optimized"
        ;;
esac

echo ""
echo "📋 训练配置:"
echo "   模型: $MODEL"
echo "   图像尺寸: ${IMGSZ}px"
echo "   训练轮数: $EPOCHS"
echo "   批次大小: $BATCH"
echo "   实验名称: $NAME"
echo ""

read -p "是否启用类别权重? (y/n): " use_weights
if [[ $use_weights == "y" || $use_weights == "Y" ]]; then
    WEIGHTS_FLAG="--use-class-weights"
    echo "✅ 启用类别权重 (Scallop加权)"
else
    WEIGHTS_FLAG=""
    echo "➖ 不使用类别权重"
fi

echo ""
echo "======================================================================"
echo "🏋️ 开始训练..."
echo "======================================================================"
echo ""

# 运行训练
python train_urpc2020_optimized.py \
    --model $MODEL \
    --epochs $EPOCHS \
    --batch $BATCH \
    --imgsz $IMGSZ \
    --name $NAME \
    $WEIGHTS_FLAG \
    --exist-ok

echo ""
echo "======================================================================"
echo "✅ 训练完成"
echo "======================================================================"
echo ""
echo "📊 查看结果: runs/train/$NAME/"
echo ""
echo "🔬 对比性能:"
echo "   python compare_models.py --optimized runs/train/$NAME/weights/best.pt"
echo ""

